import numpy as np
import matplotlib.pyplot as plt
import h5py

# 系统参数配置
class Config:
    RQ = 76.7   # 电阻-品质因数乘积
    Ql = 8893   # 外部品质因数
    beta = 1.2
    Fs = 117.402e6 / 14  # 采样率 [Hz]
    Ts = 1 / Fs  # 采样时间间隔 [s]

# 数据源配置
DATA_FILE = r'E:\MLdata\SimData\data.h5'

def read_hdf5_dataset(file_path, set_index):
    """
    从HDF5文件读取指定数据集的数据
    
    Args:
        file_path (str): HDF5文件路径
        set_index (int): 数据集索引
    
    Returns:
        tuple: 包含ABCD参数、电压波形和时间数据的元组
    """
    with h5py.File(file_path, 'r') as f:
        set_name = f'set_{set_index}'
        if set_name not in f:
            raise ValueError(f"数据集 {set_name} 不存在")
            
        dataset = f[set_name]
        
        # 读取ABCD参数（单个复数值）
        abcd_params = {
            "a": float(dataset['a_real'][0]) + 1j * float(dataset['a_imag'][0]),
            "b": float(dataset['b_real'][0]) + 1j * float(dataset['b_imag'][0]),
            "c": float(dataset['c_real'][0]) + 1j * float(dataset['c_imag'][0]),
            "d": float(dataset['d_real'][0]) + 1j * float(dataset['d_imag'][0])
        }
        
        # 读取电压波形
        voltages = {
            "Vcav": dataset['Vc_real'][:] + 1j * dataset['Vc_imag'][:],
            "Vfor": dataset['Vf_real'][:] + 1j * dataset['Vf_imag'][:],
            "Vrefl": dataset['Vr_real'][:] + 1j * dataset['Vr_imag'][:]
        }
        
        time = f['time'][:]
        
    return abcd_params, voltages, time

def print_abcd_parameters(params):
    """
    打印ABCD参数的信息
    
    Args:
        params (dict): 包含a,b,c,d参数的字典
    """
    print("\nABCD传输矩阵参数:")
    for param_name, value in params.items():
        mag = np.abs(value)
        phase = np.angle(value, deg=True)
        print(f"\n参数 {param_name}:")
        print(f"幅值: {mag:.4f}")
        print(f"相位: {phase:.2f}°")
        print(f"实部: {value.real:.4f}")
        print(f"虚部: {value.imag:.4f}")

def plot_voltage_waveforms(t_us, voltages):
    """
    绘制电压波形
    
    Args:
        t_us (array): 时间数组（微秒）
        voltages (dict): 包含电压波形的字典
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))
    
    # 绘制幅值
    ax1.plot(t_us, np.abs(voltages['Vcav']), label='|Vcav|', linewidth=1.5)
    ax1.plot(t_us, np.abs(voltages['Vfor']), label='|Vfor|', linewidth=1.5)
    ax1.plot(t_us, np.abs(voltages['Vrefl']), label='|Vrefl|', linewidth=1.5)
    ax1.set_title('Voltage Magnitudes', fontsize=12)
    ax1.grid(True, linestyle='--', linewidth=0.7)
    ax1.legend(fontsize=10)
    ax1.set_xlabel(r'Time [$\mu s$]', fontsize=10)
    ax1.set_ylabel('Magnitude', fontsize=10)
    
    # 绘制相位
    ax2.plot(t_us, np.angle(voltages['Vcav'], deg=True), label='∠Vcav', linewidth=1.5)
    ax2.plot(t_us, np.angle(voltages['Vfor'], deg=True), label='∠Vfor', linewidth=1.5)
    ax2.plot(t_us, np.angle(voltages['Vrefl'], deg=True), label='∠Vrefl', linewidth=1.5)
    ax2.set_title('Voltage Phases', fontsize=12)
    ax2.grid(True, linestyle='--', linewidth=0.7)
    ax2.legend(fontsize=10)
    ax2.set_xlabel(r'Time [$\mu s$]', fontsize=10)
    ax2.set_ylabel('Phase [deg]', fontsize=10)
    
    # 设置格式
    for ax in [ax1, ax2]:
        ax.tick_params(labelsize=10)
        for spine in ax.spines.values():
            spine.set_linewidth(1.5)
    
    plt.tight_layout()
    plt.show()

def process_file(set_index=0):
    """处理指定索引的数据集"""
    try:
        # 读取数据
        abcd_params, voltages, time = read_hdf5_dataset(DATA_FILE, set_index)
        
        # 转换时间为微秒
        t_us = time * 1e6
        
        # 打印ABCD参数信息
        print_abcd_parameters(abcd_params)
        
        # 绘制电压波形
        plot_voltage_waveforms(t_us, voltages)
        
    except Exception as e:
        print(f"处理数据集 {set_index} 时出错: {e}")

def main():
    """主函数"""
    # 处理第一个数据集
    process_file(0)

if __name__ == '__main__':
    main() 